#include "include/mm.h"
#include "include/string.h"
#include "include/print.h"
#include "include/sched.h"

unsigned long mem_size = 0; // Maximum memory size
uint8_t pages[MAX_PAGES]; // Page occupancy array

void mm_init() {
  struct e820map* e820 = (struct e820map*)0x13000; // defined in boot16.S

  // print(e820->nr_entry)
  // Traverse e820 records to find the maximum physical memory
  for (int i = 0; i < e820->nr_entry; i++) {
    // print(e820->map[i].addr)
    // print(e820->map[i].size)
    // print(e820->map[i].type)
    if (e820->map[i].type == E820_RAM) {
      unsigned long tmp = e820->map[i].addr + e820->map[i].size; // end address = start + length
      if (tmp > mem_size) {
        mem_size = tmp;
      }
    }
  }

  print(mem_size);

  // Initialize all page occupancy to 0
  memset(pages, 0, MAX_PAGES);

  // Mark pages that have been occupied by the kernel as 1
  for (int i = 0; i < KERNEL_PAGE_NUM; i++) {
    pages[i] = 1;
  }

  // Establish a linear mapping from full memory physical address to virtual address
  map_range(TASK0_PML4, (unsigned long)VA(0), 0, 0, (mem_size + PAGE_SIZE - 1) / PAGE_SIZE);
}

/**
 * Page Table Mapping Establishment Function
 * - from_va and to_pa move back n 4KB(page size) at the same pace
 *
 * @param pml4_pa The physical base address of the fourth level root page table
 * @param from_va The starting virtual address that needs to be mapped
 * @param to_pa   The starting physical address that needs to be mapped
 * @param us      Access permissions for page tables
 *                0x4(0100): Memory mapped to page table entries allows application access
 *                0        : Only for kernel
 * @param npage   Number of mapped pages
 */
void map_range(unsigned long pml4_pa, unsigned long from_va, unsigned long to_pa, char us, long npage) {
  long n = 0;

  while (n < npage) {
    // pml4
    unsigned long* page_va = VA(pml4_pa); // physical addr -> virtual addr
    short index = (from_va >> 39) & 0x1ff; // Parse the 4th level index, taking 47-39 bits for a total of 9 bits
    // Parsing page table entry at the index of the 4th level page table,
    // map the physical base address of the 3rd level page table
    unsigned long entry = *(page_va + index);
    if (!(entry & 0x1)) { // Determine whether a page exists by the presence bit of page table entry
      // Assign a new page when it does not exist, and set permissions
      *(page_va + index) = alloc_page() | 0x3 | us; // Write back to page table entry
      entry = *(page_va + index);
    }

    // pml3
    page_va = VA(entry & 0xffffffffff000); // physical -> virtual, 4KB address alignment, using bottom 3 bits of hexadecimal
    index = (from_va >> 30) & 0x1ff; // Parse the 3th level index, taking 38-30 bits for a total of 9 bits
    entry = *(page_va + index); // physical base address of the 2rd level page table
    if (!(entry & 0x1)) {
      *(page_va + index) = alloc_page() | 0x3 | us;
      entry = *(page_va + index);
    }

    // pml2
    page_va = VA(entry & 0xffffffffff000); // physical -> virtual
    index = (from_va >> 21) & 0x1ff; // Parse the 2th level index, taking 29-21 bits for a total of 9 bits
    entry = *(page_va + index); // physical base address of the 1st level page table
    if (!(entry & 0x1)) {
      *(page_va + index) = alloc_page() | 0x3 | us;
      entry = *(page_va + index);
    }

    // pml1
    page_va = VA(entry & 0xffffffffff000); // physical -> virtual
    index = (from_va >> 12) & 0x1ff; // Parse the 1th level index, taking 20-12 bits for a total of 9 bits
    entry = *(page_va + index); // physical base address of the physical page
    if (!(entry & 0x1)) {
      // Write physical page address back to page table entry
      *(page_va + index) = (to_pa + PAGE_SIZE * n) | 0x3 | us; // to_pa move back n 4KB
    }

    n++;
    from_va += PAGE_SIZE; // from_va move back n 4KB
  }
}

/**
 * page fault interrupt handling function
 * 
 * @param addr Missing page virtual address
 */
void do_page_fault(unsigned long addr) {
  unsigned long pa = alloc_page(); // page physical address
  map_range(current->pml4, addr, pa, 0x4, 1); // mapping 1 page
}